import vitis
import os
import sys
import shutil
from pathlib import Path

def run_hls(syn_file: str, tb_file: str = None, header_file: str = None, top_name: str = None, config_file: str = 'hls_config.cfg'):
    """
    Run HLS using Vitis HLS Python API.
    
    Args:
        syn_file (str): Path to synthesis file
        tb_file (str, optional): Path to testbench file
        header_file (str, optional): Path to header file
        top_name (str): Name of the top function
        config_file (str): Path to configuration file
    """
    if not top_name:
        raise ValueError("Top function name must be provided")
        
    # Get current working directory
    cwd = os.getcwd()
    print(f"Current working directory: {cwd}")
    
    # Convert relative paths to absolute paths
    syn_file = os.path.abspath(syn_file)
    if tb_file:
        tb_file = os.path.abspath(tb_file)
    if header_file:
        header_file = os.path.abspath(header_file)
    
    print(f"Synthesis file: {syn_file} (exists: {os.path.exists(syn_file)})")
    if tb_file:
        print(f"Testbench file: {tb_file} (exists: {os.path.exists(tb_file)})")
    if header_file:
        print(f"Header file: {header_file} (exists: {os.path.exists(header_file)})")
    
    # Clean up workspace directory
    workspace_dir = os.path.join(cwd, 'workspace')
    if os.path.exists(workspace_dir):
        try:
            shutil.rmtree(workspace_dir)
            print(f"Cleaned up workspace directory: {workspace_dir}")
        except Exception as e:
            print(f"Warning: Could not clean up workspace directory: {str(e)}")
    
    # Initialize session
    client = vitis.create_client()
    client.set_workspace(path='./workspace')
    
    # Component name based on top function
    comp_name = top_name
    
    # Create component
    comp = client.create_hls_component(
        name=comp_name,
        cfg_file=[config_file],
        template='empty_hls_component'
    )
    
    # Wait for component directory to be created
    comp_dir = os.path.abspath(f'./workspace/{comp_name}')
    max_retries = 10
    retry_delay = 1  # seconds
    
    for retry in range(max_retries):
        if os.path.exists(comp_dir):
            break
        import time
        time.sleep(retry_delay)
    else:
        raise RuntimeError(f"Component directory {comp_dir} was not created")
    
    print(f"Component directory: {comp_dir} (exists: {os.path.exists(comp_dir)})")
    
    # Copy source files to component directory
    try:
        # Create target files
        syn_target = os.path.join(comp_dir, os.path.basename(syn_file))
        print(f"Creating synthesis file: {syn_target}")
        with open(syn_file, 'r', encoding='utf-8') as src, open(syn_target, 'w', encoding='utf-8') as dst:
            dst.write(src.read())
        
        if tb_file:
            tb_target = os.path.join(comp_dir, os.path.basename(tb_file))
            print(f"Creating testbench file: {tb_target}")
            with open(tb_file, 'r', encoding='utf-8') as src, open(tb_target, 'w', encoding='utf-8') as dst:
                dst.write(src.read())
        
        if header_file:
            header_target = os.path.join(comp_dir, os.path.basename(header_file))
            print(f"Creating header file: {header_target}")
            with open(header_file, 'r', encoding='utf-8') as src, open(header_target, 'w', encoding='utf-8') as dst:
                dst.write(src.read())
    except Exception as e:
        print(f"Error copying files: {str(e)}")
        print(f"Source file: {syn_file}")
        print(f"Target directory: {comp_dir}")
        print(f"Current working directory: {os.getcwd()}")
        print(f"File exists: {os.path.exists(syn_file)}")
        if tb_file:
            print(f"Testbench exists: {os.path.exists(tb_file)}")
        if header_file:
            print(f"Header exists: {os.path.exists(header_file)}")
        raise
    
    # Get handle of config file and set options
    cfg_file = client.get_config_file(path=f'./workspace/{comp_name}/{config_file}')
    
    # Set basic configuration
    cfg_file.set_value(key='part', value='xc7z020clg484-1')  # Default Zynq-7000 part
    
    # Set HLS specific configuration
    # Use relative paths for synthesis and testbench files
    cfg_file.set_value(section='hls', key='syn.file', value=os.path.basename(syn_file))
    if tb_file:
        cfg_file.set_values(section='hls', key='tb.file', values=[os.path.basename(tb_file)])
    
    # Set top function name
    print(f"Setting top function to: {top_name}")
    cfg_file.set_value(section='hls', key='syn.top', value=top_name)
    
    # Set other HLS configurations
    cfg_file.set_value(section='hls', key='clock', value='10')  # 100MHz
    cfg_file.set_value(section='hls', key='flow_target', value='vivado')
    cfg_file.set_value(section='hls', key='package.output.format', value='rtl')
    cfg_file.set_value(section='hls', key='csim.code_analyzer', value='0')
    
    # Get component handle and run flow steps
    comp = client.get_component(name=comp_name)
    
    # Run simulation if testbench is provided
    if tb_file:
        comp.run(operation='C_SIMULATION')
    
    # Run synthesis
    comp.run(operation='SYNTHESIS')
    
    # # Run co-simulation if testbench is provided
    # if tb_file:
    #     comp.run(operation='CO_SIMULATION')

if __name__ == "__main__":
    # Get arguments from command line
    if len(sys.argv) < 2:
        print("Usage: python run.py <syn_file> [tb_file] [header_file] [top_name] [config_file]")
        sys.exit(1)
    
    syn_file = sys.argv[1]
    tb_file = sys.argv[2] if len(sys.argv) > 2 and sys.argv[2] != 'None' else None
    header_file = sys.argv[3] if len(sys.argv) > 3 and sys.argv[3] != 'None' else None
    top_name = sys.argv[4] if len(sys.argv) > 4 else None
    config_file = sys.argv[5] if len(sys.argv) > 5 else 'hls_config.cfg'
    
    if not top_name:
        print("Error: Top function name must be provided")
        sys.exit(1)
    
    run_hls(syn_file, tb_file, header_file, top_name, config_file) 